# 雷电预测数据集类
from typing import Dict, Tuple, List
import os
from datetime import datetime
import random
import numpy as np
import matplotlib.pyplot as plt
import torch
from torch.utils.data import Dataset
from apps.tlp.tlp_config import TlpConfig as TG

class TlpDs(Dataset):
    def __init__(self):
        super(TlpDs, self).__init__()
        self.X, self.y = None, None
        self.ds_size = 0
        print(f'计算数据集大小...')
        self.ds_folder = 'work/tlp/datasets'
        for root, dirs, fns in os.walk(self.ds_folder):
            for fn in fns:
                if 'X_' in fn:
                    ffn = f'{root}/{fn}'
                    X = torch.load(ffn, weights_only=True)
                    self.ds_size += X.shape[0]
        self.curr_seq = 0
        X_ffn= f'{self.ds_folder}/X_{self.curr_seq:05d}.pt'
        self.X = torch.load(X_ffn, weights_only=True)
        y_ffn= f'{self.ds_folder}/y_{self.curr_seq:05d}.pt'
        self.y = torch.load(y_ffn, weights_only=True)
        
    def __len__(self):
        return self.ds_size

    def __getitem__(self, index: int):
        seq = int(index / TG.BLOCK_SIZE)
        idx = index % TG.BLOCK_SIZE
        if self.curr_seq != seq:
            self.curr_seq = seq
            X_ffn= f'{self.ds_folder}/X_{self.curr_seq:05d}.pt'
            self.X = torch.load(X_ffn, weights_only=True)
            y_ffn= f'{self.ds_folder}/y_{self.curr_seq:05d}.pt'
            self.y = torch.load(y_ffn, weights_only=True)
        return self.X[idx], self.y[idx]


    IDX_id = 0 # 为预警记录的唯一编号，注意 1208, 1589, 4878, 13017, 21759, 32426, 33790 由于缺少预警开始或结束时间，对应行的数据已删除；
    IDX_did = 1 # 设备识别码；
    IDX_dname = 2 # 设备名称；
    IDX_device_lng = 3 # 设备所在经度；
    IDX_device_lat = 4 # 设备所在纬度；
    IDX_max_elevel = 5 # 预警最大等级（1-5）;
    IDX_warning_start_time = 6 # 预警开始时间；
    IDX_warning_end_time = 7 # 预警结束时间；
    IDX_first_lightning_time = 8 # 预警持续时间内首次闪电发生时间；
    IDX_first_lightning_lng = 9 # 首次闪电发生经度；
    IDX_first_lightning_lat = 10 # 首次闪电发生纬度；
    IDX_first_lightning_distance = 11 # 首次闪电距设备的距离（km）；
    IDX_first_lightning_height = 12 # 首次闪电高度*********************************************
    IDX_first_lightning_strength = 13 # 首次闪电发生强度 （kA）；
    IDX_first_lw_level = 14 # 首次闪电发生时对应的设备预警等级；
    IDX_last_lightning_time = 15 # 预警持续时间内末次闪电发生时间；
    IDX_last_lightning_lng = 16 # 末次闪电发生经度；
    IDX_last_lightning_lat = 17 # 末次闪电发生纬度；
    IDX_last_lightning_distance = 18 # 末次闪电距设备的距离（km）；
    IDX_last_lightning_height = 19
    IDX_last_lightning_strength = 20 # 末次闪电发生强度 （kA）；
    IDX_last_lw_level = 21 # 末次闪电发生时对应的设备预警等级。  

    @staticmethod
    def startup(params:Dict = {}) -> None:
        # 设置中文字体
        plt.rcParams['font.sans-serif'] = ['SimHei']  # 使用黑体
        plt.rcParams['axes.unicode_minus'] = False  # 解决负号显示问题
        rec_dict = TlpDs.load_result_csv(csv_fn='work/tlp/datas/warnings1-8.csv')

    @staticmethod
    def process_raw_datas(txt_fn:str, rec_dict:Dict) -> None:
        X_, y_ = None, None
        with open(txt_fn, 'r', encoding='utf-8') as rfd:
            num = 0
            for row in rfd:
                arrs0 = row.strip().split(' ')
                if len(arrs0) <= 1:
                    continue
                row_id = arrs0[0]
                vals = arrs0[1].split(',')
                Xs, ys = TlpDs.extract_frames(raw_datas=vals, rec=rec_dict[row_id])
                if Xs is None or ys is None:
                    continue
                if X_ is None:
                    X_ = Xs
                    y_ = ys
                else:
                    X_ = torch.vstack((X_, Xs))
                    y_ = torch.vstack((y_, ys))
                if X_.shape[0] > TG.BLOCK_SIZE:
                    xi = X_[:TG.BLOCK_SIZE] # 存入文件
                    yi = y_[:TG.BLOCK_SIZE]
                    x_fn = f'./work/tlp/datasets/X_{TG.SEQ:05d}.pt'
                    y_fn = f'./work/tlp/datasets/y_{TG.SEQ:05d}.pt'
                    torch.save(xi, x_fn)
                    torch.save(yi, y_fn)
                    X_ = X_[TG.BLOCK_SIZE:]
                    y_ = y_[TG.BLOCK_SIZE:]
                    TG.SEQ += 1
                num += 1
                if num % 1000 == 0:
                    print(f'处理完成{num}行...')

    @staticmethod
    def extract_frames(raw_datas:List, rec:object) -> Tuple[torch.Tensor, torch.Tensor]:
        Xs = None
        ys = None
        cnt = len(raw_datas)
        for idx in range(0, cnt-TG.FRAME_SIZE, TG.FRAME_STEP):
            v_x = torch.tensor([float(item) for item in raw_datas[idx : idx+TG.FRAME_SIZE]])
            y = TlpDs.classify_frame(idx, idx+TG.FRAME_STEP, rec['tl_start_idx'], rec['tl_end_idx'])
            if Xs is None:
                Xs = v_x
            else:
                Xs = torch.vstack((Xs, v_x))
            if ys is None:
                ys = torch.tensor([y])
            else:
                ys = torch.vstack((ys, torch.tensor([y])))
        return Xs, ys

    @staticmethod
    def classify_frame(start_idx:int, end_idx:int, tl_start_idx:int, tl_end_idx:int) -> int:
        '''
        start_idx: 帧开始的下标
        end_idx: 帧结束的下标
        tl_start_idx: 雷电开始的下标
        tl_end_idx: 雷电结束的下标
        返回值：
          当tl_start_idx=-1时，返回0-正常情况；
          当start_idx < tl_start_idx and start_idx + 10*60 > tl_start_idx: 返回3-红色预警
          当start_idx < tl_start_idx and start_idx + 20*60 > tl_start_idx: 返回2-橙色预警
          当start_idx < tl_start_idx and start_idx + 30*60 > tl_start_idx: 返回1-黄色预警
          当(start_idx>=tl_start_idx and start_idx<=tl_end_idx) or (end_idx>=tl_start_idx and end_idx<=tl_end_idx): 返回4-雷电中
        '''
        if tl_start_idx == -1:
            return 0
        if end_idx < tl_start_idx and end_idx + TG.WARN_YELLOW < tl_start_idx:
            return 0
        if end_idx < tl_start_idx and end_idx + TG.WARN_RED > tl_start_idx:
            return 3
        if end_idx < tl_start_idx and end_idx + TG.WARN_BROWN > tl_start_idx:
            return 2
        if end_idx < tl_start_idx and end_idx + TG.WARN_YELLOW > tl_start_idx:
            return 1
        if (start_idx>=tl_start_idx and start_idx<=tl_end_idx) or (end_idx>=tl_start_idx and end_idx<=tl_end_idx):
            return 4
        return 0

    @staticmethod
    def seconds_between(t1:str, t2:str) -> int:
        '''
        求出t1和t2间间隔的秒数，时间格式为：'yyyy-mm-dd hh:mm:ss'
        '''
        # 将字符串转换为datetime对象
        time_format = '%Y-%m-%d %H:%M:%S'
        dt1 = datetime.strptime(t1, time_format)
        dt2 = datetime.strptime(t2, time_format)
        # 计算两个时间之间的差异，并转换为秒
        time_difference = dt2 - dt1
        return int(time_difference.total_seconds())

    @staticmethod
    def load_result_csv(csv_fn:str) -> Dict:
        '''
        从结果csv文件中，读出每条记录闪电开始下标值和结束下标值
        '''
        rec_dict = {}
        with open(csv_fn, 'r', encoding='utf-8') as rfd:
            first_run = True
            num = 0
            for row in rfd:
                if first_run:
                    first_run = False
                    continue
                row = row.strip()
                arrs0 = row.split(',')
                if arrs0[TlpDs.IDX_first_lightning_time] == '' or arrs0[TlpDs.IDX_first_lightning_time] is None:
                    tl_start_idx = -1
                else:
                    tl_start_idx = 10*60 + TlpDs.seconds_between(arrs0[TlpDs.IDX_warning_start_time][1:-1], arrs0[TlpDs.IDX_first_lightning_time][1:-1])
                if arrs0[TlpDs.IDX_last_lightning_time] == '' or arrs0[TlpDs.IDX_last_lightning_time] is None:
                    tl_end_idx = -1
                else:
                    tl_end_idx = 10*60 + TlpDs.seconds_between(arrs0[TlpDs.IDX_warning_start_time][1:-1], arrs0[TlpDs.IDX_last_lightning_time][1:-1])
                rec_dict[arrs0[TlpDs.IDX_id]] = {
                    'warn_time': arrs0[TlpDs.IDX_warning_start_time],
                    'tl_start_time': arrs0[TlpDs.IDX_first_lightning_time],
                    'tl_start_idx': tl_start_idx,
                    'tl_end_time': arrs0[TlpDs.IDX_last_lightning_time],
                    'tl_end_idx': tl_end_idx
                }
                num += 1
                if num % 1000 == 0:
                    print(f'处理完成{num}条记录...')
        return rec_dict

    @staticmethod
    def data_visualize(params:Dict = {}) -> None:
        normal_rows, normal_ids, tl_rows, tl_ids = TlpDs.load_tl_datas(params=params)
        # 读入所有数据
        datas = TlpDs.load_datas(params=params)
        # fig, ax = plt.subplots(2, 5, sharex=True, sharey=True, figsize=(7, 2), dpi=300)
        # plt.subplots_adjust(wspace=0.1, hspace=0.05)
        # 处理正常情况
        nids = list(normal_ids)
        normal_recs = [] # 正常情况下的变化数据
        normal_num = 5 # 取5个正常情况
        for idx in range(normal_num):
            rv = random.randint(0, len(nids))
            normal_recs.append(nids[rv])
        for idx in range(normal_num):
            ny = datas[normal_recs[idx]]
            nx = np.arange(ny.shape[0])
            # ax[0][idx].plot(nx1, ny1)
            # ax[0][idx].set_title(f'正常情况_{idx}')
            plt.plot(nx, ny)
            plt.title(f'正常情况_{idx}')
            plt.show()
        # 处理雷电情况
        tids = list(tl_ids)
        tl_num = 5
        tl_recs = [] # 雷电情况下的变化数据
        for idx in range(tl_num):
            rv = random.randint(0, len(tids))
            tl_recs.append(tids[rv])
        for idx in range(tl_num):
            ty = datas[tl_recs[idx]]
            tx = np.arange(ty.shape[0])
            # ax[1][idx].plot(tx, ty)
            # ax[1][idx].set_title(f'雷电情况_{idx}')
            plt.plot(tx, ty)
            plt.title(f'雷电情况_{idx}')
            plt.show()
        # plt.show() 
    

    @staticmethod
    def load_tl_datas(params: Dict = {}) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
        '''
        从warnings1-8.csv中读出雷电记录和非雷电记录
        返回值：
            normal_rows: 正常情况，即预报发生雷电而实际未发生雷电
            normal_ids：记录编号（与electric.txt）中记录相对应
            tl_rows：雷电记录，即预报发生雷电而实际发生雷电
            tl_ids：记录编号（与electric.txt）中记录相对应
        '''
        normal_rows = []
        normal_ids = set()
        tl_rows = []
        tl_ids = set()
        # 读取所有记录
        csv_fn = 'work/tlp/datas/warnings1-8.csv'
        first_row = True
        with open(csv_fn, 'r', encoding='utf-8') as cfd:
            for row in cfd:
                if first_row:
                    first_row = False
                    continue
                arrs = row.strip().split(',')
                item = {
                    'id': arrs[TlpDs.IDX_id], # 为预警记录的唯一编号，注意 1208, 1589, 4878, 13017, 21759, 32426, 33790 由于缺少预警开始或结束时间，对应行的数据已删除；
                    'did': arrs[TlpDs.IDX_did], # 设备识别码；
                    'dname': arrs[TlpDs.IDX_dname], # 设备名称；
                    'device_lng': arrs[TlpDs.IDX_device_lng], # 设备所在经度；
                    'device_lat': arrs[TlpDs.IDX_device_lat], # 设备所在纬度；
                    'max_elevel': arrs[TlpDs.IDX_max_elevel], # 预警最大等级（1-5）;
                    'warning_start_time': arrs[TlpDs.IDX_warning_start_time], # 预警开始时间；
                    'warning_end_time': arrs[TlpDs.IDX_warning_end_time], # 预警结束时间；
                    'first_lightning_time': arrs[TlpDs.IDX_first_lightning_time], # 预警持续时间内首次闪电发生时间；
                    'first_lightning_lng': arrs[TlpDs.IDX_first_lightning_lng], # 首次闪电发生经度；
                    'first_lightning_lat': arrs[TlpDs.IDX_first_lightning_lat], # 首次闪电发生纬度；
                    'first_lightning_distance': arrs[TlpDs.IDX_first_lightning_distance], # 首次闪电距设备的距离（km）；
                    'first_lightning_strength': arrs[TlpDs.IDX_first_lightning_strength], # 首次闪电发生强度 （kA）；
                    'first_lw_level': arrs[TlpDs.IDX_first_lw_level], # 首次闪电发生时对应的设备预警等级；
                    'last_lightning_time': arrs[TlpDs.IDX_last_lightning_time], # 预警持续时间内末次闪电发生时间；
                    'last_lightning_lng': arrs[TlpDs.IDX_last_lightning_lng], # 末次闪电发生经度；
                    'last_lightning_lat': arrs[TlpDs.IDX_last_lightning_lat], # 末次闪电发生纬度；
                    'last_lightning_distance': arrs[TlpDs.IDX_last_lightning_distance], # 末次闪电距设备的距离（km）；
                    'last_lightning_strength': arrs[TlpDs.IDX_last_lightning_strength], # 末次闪电发生强度 （kA）；
                    'last_lw_level': arrs[TlpDs.IDX_last_lw_level], # 末次闪电发生时对应的设备预警等级。
                }
                if item['first_lightning_time'] == '':
                    normal_rows.append(item)
                    normal_ids.add(item['id'])
                else:
                    tl_rows.append(item)
                    tl_ids.add(item['id'])
        return normal_rows, normal_ids, tl_rows, tl_ids

    
    @staticmethod
    def load_datas(params:Dict = {}) -> Dict:
        data_fn = 'work/tlp/datas/electric.txt'
        datas = {} # id: np.ndarray
        datas_raw = []
        num = 0
        with open(data_fn, 'r', encoding='utf-8') as rfd:
            for row in rfd:
                arrs0 = row.strip().split(' ')
                if len(arrs0) <= 1:
                    continue
                arrs1 = arrs0[1].split(',')
                item = []
                for x in arrs1:
                    ki = float(x)
                    if ki > 110:
                        print(f'########### {ki};')
                    item.append(ki)
                datas[arrs0[0]] = np.array(item)
                num += 1
                if num % 1000 == 0:
                    print(f'处理{num}行...')
        return datas